#include "StdAfx.h"
#include "TestRunnerExchange2.h"
#include <Tlhelp32.h>

CTestRunnerExchange2::CTestRunnerExchange2(void)	:
	m_bConnected(FALSE), m_nMmfSize(0)
{
	DWORD dwDevEnvProcId = DefineDebuggerProcId();
	swprintf_s(m_wsz, _countof(m_wsz), L"cputaddin_%i", dwDevEnvProcId);
	if ((dwDevEnvProcId != 0xFFFFFFFF) &&
		(m_hTestRunnerWnd = FindTestRunnerWnd(dwDevEnvProcId)))
	{
		m_bConnected = TRUE;
		m_hEvent.Attach(CreateEvent(NULL, FALSE, FALSE, m_wsz));
		NotifyServer(stepStartNotification, GetCurrentProcessId());
	}
}

CTestRunnerExchange2::~CTestRunnerExchange2(void)
{
	Dispose();
}

void CTestRunnerExchange2::Start(EFilter eFilter)
{
}

void CTestRunnerExchange2::GroupStart(LPCWSTR lpGroupName, DWORD nGroupParam)
{
	if (IsConnected())
	{
		LPVOID lp = MapViewOfFile(m_hMmf, FILE_MAP_WRITE, 0, 0, 4);
		
		if (lp)
		{
			*((LPDWORD)lp) = nGroupParam;

			UnmapViewOfFile(lp);

			NotifyServer(stepGroupStartNotification, 4);
		}
		else
		{
			Dispose();
		}
	}
}

void CTestRunnerExchange2::TestStart(LPCWSTR /*lpGroupName*/, LPCWSTR /*lpTestName*/, DWORD nGroupParam, DWORD nTestParam)
{
	if (IsConnected())
	{
		LPVOID lp = MapViewOfFile(m_hMmf, FILE_MAP_WRITE, 0, 0, 8);
		
		if (lp)
		{
			LPDWORD lpDw = (LPDWORD)lp;
			*lpDw = nGroupParam; ++lpDw;
			*lpDw = nTestParam;

			UnmapViewOfFile(lp);

			NotifyServer(stepTestStartNotification, 8);
		}
		else
		{
			Dispose();
		}
	}
}

void CTestRunnerExchange2::TestEnd(LPCWSTR lpGroupName, LPCWSTR lpTestName, ERunStatus eRunStatus, LPCWSTR lpReport, __time64_t tDuration, DWORD nGroupParam, DWORD nTestParam)
{
	if (IsConnected())
	{
		DWORD dwLen = static_cast<DWORD>(wcslen(lpReport));
		DWORD dwSize = sizeof(eRunStatus) + sizeof(tDuration) + sizeof(nGroupParam) + sizeof(nTestParam) + sizeof(dwLen) + sizeof(wchar_t) * dwLen + sizeof(L'\0');

		if (dwSize > m_nMmfSize)
		{
			//crop the report string if the overall is too big
			dwLen -= m_nMmfSize - dwSize;
			dwSize = m_nMmfSize;
		}
				
		LPVOID lp = MapViewOfFile(m_hMmf, FILE_MAP_WRITE, 0, 0, dwSize);
		
		if (lp)
		{
			LPDWORD lpDw = (LPDWORD)lp;
			*lpDw = nGroupParam; ++lpDw;
			*lpDw = nTestParam; ++lpDw;
			*lpDw = eRunStatus; ++lpDw;
			*((__time64_t*)lpDw) = tDuration; ++lpDw; ++lpDw;
			*lpDw = dwLen + sizeof(L'\0'); ++lpDw;

			LPWSTR lpW = (LPWSTR)lpDw;
			if (dwLen)
			{
				wcsncpy_s(lpW, (((LPWSTR)((LPBYTE)lp + dwSize)) - lpW), lpReport, dwLen);
				lpW += dwLen;
			}
			*lpW = L'\0';


			UnmapViewOfFile(lp);

			NotifyServer(stepTestEndNotification, dwSize);
		}
		else
		{
			Dispose();
		}
	}
}

void CTestRunnerExchange2::GroupEnd(LPCWSTR lpGroupName, ERunStatus eRunStatus, LPCWSTR lpReport, __time64_t tDuration, DWORD nGroupParam)
{
	if (IsConnected())
	{
		DWORD dwLen = static_cast<DWORD>(wcslen(lpReport));
		DWORD dwSize = sizeof(eRunStatus) + sizeof(tDuration) + sizeof(nGroupParam) + sizeof(dwLen) + sizeof(wchar_t) * dwLen + sizeof(L'\0');

		if (dwSize > m_nMmfSize)
		{
			//crop the report string if the overall is too big
			dwLen -= m_nMmfSize - dwSize;
			dwSize = m_nMmfSize;
		}
				
		LPVOID lp = MapViewOfFile(m_hMmf, FILE_MAP_WRITE, 0, 0, dwSize);
		
		if (lp)
		{
			LPDWORD lpDw = (LPDWORD)lp;
			*lpDw = nGroupParam; ++lpDw;
			*lpDw = eRunStatus; ++lpDw;
			*((__time64_t*)lpDw) = tDuration; ++lpDw; ++lpDw;
			*lpDw = dwLen + sizeof(L'\0'); ++lpDw;

			LPWSTR lpW = (LPWSTR)lpDw;
			if (dwLen)
			{
				wcsncpy_s(lpW, (((LPWSTR)((LPBYTE)lp + dwSize)) - lpW), lpReport, dwLen);
				lpW += dwLen;
			}
			*lpW = L'\0';

			UnmapViewOfFile(lp);

			NotifyServer(stepGroupEndNotification, dwSize);
		}
		else
		{
			Dispose();
		}
	}
}

void CTestRunnerExchange2::End()
{
	if (IsConnected())
	{
		NotifyServer(stepEndNotification);
	}

	Dispose();
}

void CTestRunnerExchange2::Dispose()
{
	m_bConnected = FALSE;
	m_nMmfSize = 0;
	m_hEvent.Close();
	m_hMmf.Close();
}

BOOL CTestRunnerExchange2::NotifyServer(ETestRunExchangeStep step, LPARAM lParam /*= 0*/)
{
	if ((FALSE == PostMessage(m_hTestRunnerWnd, WM_TESTAPPNOTIFY, step, lParam)) ||
		(WAIT_OBJECT_0 != WaitForSingleObject(m_hEvent, c_nTimeout)))
	{
		Dispose();
		return FALSE;
	}
	return TRUE;
}

BOOL CTestRunnerExchange2::DefineTestsToRun(TestPackTransData & rTestPackTransData)
{
	BOOL bRet = FALSE;
	if (IsConnected())
	{
		wchar_t wsz[40];
		swprintf_s(wsz, _countof(wsz), L"%s_%i", m_wsz, GetCurrentProcessId());

		m_nMmfSize = rTestPackTransData.DefineSize();
		SYSTEM_INFO sSI;
		GetSystemInfo(&sSI);
		if (m_nMmfSize < sSI.dwAllocationGranularity)
			m_nMmfSize = sSI.dwAllocationGranularity;

		m_hMmf.Attach(CreateFileMappingW(INVALID_HANDLE_VALUE, NULL, PAGE_READWRITE, 0, m_nMmfSize, wsz));

		if ((m_hMmf.m_h) &&
			(m_hMmf.m_h != INVALID_HANDLE_VALUE))
		{
			LPVOID lp = MapViewOfFile(m_hMmf, FILE_MAP_WRITE, 0, 0, m_nMmfSize);
			rTestPackTransData.Serialize((LPBYTE)lp, m_nMmfSize);
			UnmapViewOfFile(lp);

			if (NotifyServer(stepTestSet, m_nMmfSize))
			{
				lp = MapViewOfFile(m_hMmf, FILE_MAP_WRITE, 0, 0, m_nMmfSize);
				rTestPackTransData.DeSerialize((LPBYTE)lp, m_nMmfSize);
				UnmapViewOfFile(lp);

				bRet = TRUE;
			}
		}
		else
		{
			Dispose();
		}
	}
	
	return bRet;
}

DWORD CTestRunnerExchange2::DefineDebuggerProcId()
{
	DWORD dwRet = 0xFFFFFFFF;
	DWORD dwCurProcId = GetCurrentProcessId();
	CHandle csSnapshot(CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0));
	PROCESSENTRY32 sPE;
	sPE.dwSize = sizeof(sPE);
	BOOL bContinue = TRUE;

	//In case Running (Ctrl + F5) immediate parent would be cmd.exe
	//And if the launch platform x64 for console parent would be msvsmon.exe

	while ((bContinue) && (Process32First(csSnapshot, &sPE)))
	{
		bContinue = FALSE;
		do
		{
			if (dwCurProcId == sPE.th32ProcessID)
			{
				CPathW path(sPE.szExeFile);
				int nFileIndx = path.FindFileName();
				int nExtIndx = path.FindExtension();
				CAtlStringW str = path.m_strPath.Mid(nFileIndx, nExtIndx - nFileIndx);
				if ((nFileIndx != -1) && (nExtIndx != -1) && 
					(0 == path.m_strPath.Mid(nFileIndx, nExtIndx - nFileIndx).CollateNoCase(L"devenv")))
				{
					dwRet = dwCurProcId;
				} else {
					dwCurProcId = sPE.th32ParentProcessID;
					bContinue = TRUE;
				}
				break;
			}
		}
		while (Process32Next(csSnapshot, &sPE));
	}

	return dwRet;
}

BOOL CALLBACK CTestRunnerExchange2::EnumWindowsProc(HWND hwnd, LPARAM lParam)
{
	DWORD dwProcId;
	SEnumWndParams* lp = (SEnumWndParams*)lParam;
	GetWindowThreadProcessId(hwnd, &dwProcId);
	if (lp->dwProcId == dwProcId)
	{
		FindTestRunnerWnd(hwnd, lp);
	}

	return lp->hWnd == NULL;
}

HWND CTestRunnerExchange2::FindTestRunnerWnd(DWORD dwProcId)
{
	SEnumWndParams sParams;
	sParams.dwProcId = dwProcId;
	sParams.hWnd = NULL;
	EnumWindows(EnumWindowsProc, (LPARAM)&sParams);

	return sParams.hWnd;
}

BOOL CALLBACK CTestRunnerExchange2::EnumChildProc(HWND hwnd, LPARAM lParam)
{
	SEnumWndParams * pParams = (SEnumWndParams*)lParam;
	FindTestRunnerWnd(hwnd, pParams);

	return pParams->hWnd == NULL;
}

void CTestRunnerExchange2::FindTestRunnerWnd(HWND hwnd, SEnumWndParams * pParams)
{
	HWND hFound = FindWindowExW(hwnd, NULL, NULL, L"Test Runner");
	if (hFound)
	{
		HWND hChild = GetWindow(hFound, GW_CHILD);
		LONG nUserData = GetWindowLong(hChild, GWL_USERDATA);

		if (nUserData == pParams->dwProcId)
		{
			pParams->hWnd = hChild;
		}
	} 
	else
	{
		EnumChildWindows(hwnd, EnumChildProc, (LPARAM)pParams);
	}
}